use std::cell::RefCell;
use std::rc::Rc;
use sha2::{Digest, Sha256};

pub struct MerkleTree {
    pub root: Rc<MerkleNode>,
}

#[derive(Clone)]
pub struct MerkleNode {
    left: Option<Rc<MerkleNode>>,
    right: Option<Rc<MerkleNode>>,
    pub(crate) data: Vec<u8>,
}

impl MerkleNode {
    pub fn new_merkle_node(left: Option<Rc<MerkleNode>>, right: Option<Rc<MerkleNode>>, data: &Vec<u8>) -> Rc<MerkleNode> {
        let mut hasher = Sha256::new();
        let prev_hashes: Vec<u8>;

        if left.is_none() && right.is_none() {
            hasher.update(data);
            prev_hashes = hasher.finalize().to_vec();

            Rc::new(MerkleNode {
                left: None,
                right: None,
                data: prev_hashes,
            })
        } else {
            let lf = left.unwrap();
            let rg = right.unwrap();
            prev_hashes = [lf.data.clone(), rg.data.clone()].concat();
            Rc::new(MerkleNode {
                left: Some(Rc::clone(&lf)),
                right: Some(Rc::clone(&rg)),
                data: prev_hashes,
            })
        }
    }
}

impl MerkleTree {
    pub fn new_merkle_tree(mut data: Vec<Vec<u8>>) -> MerkleTree {
        let mut nodes: Vec<Rc<MerkleNode>> = vec![];

        if data.len() % 2 != 0 {
            data.push(data[data.len() - 1].clone());
        }

        for datum in &data {
            let node = MerkleNode::new_merkle_node(None, None, datum);
            nodes.push(Rc::clone(&node));
        }

        for _i in 0..data.len() / 2 {
            let mut new_level: Vec<Rc<MerkleNode>> = vec![];

            let mut j = 0;
            while j < nodes.len() {
                let node = MerkleNode::new_merkle_node(Some(Rc::clone(&nodes[j])), Some(Rc::clone(&nodes[j + 1])), &vec![]);
                new_level.push(node);
                j += 2
            }

            nodes = new_level;
        }

        let m_tree = MerkleTree { root: Rc::clone(&nodes[0]) };

        m_tree
    }
}